#!/usr/bin/env python
#
# Upload a file, directory, or .torrent+file(s) to Google App Engine datastore,
# organized in a bittorrent- and datastore- friendly format (files are split
# into fixed-size chunks).
# Can also delete stuff you uploaded before.
#
# Uses Google's remote_api to do the transfer; you must have something like
# this:
# handlers:
# - url: /remote_api
#   script: $PYTHON_LIB/google/appengine/ext/remote_api/handler.py
#   login: admin
# ... in your app.yaml file.
#
# Run with no arguments to get usage.
#

import commands
import getpass
import hashlib
import logging
import optparse
import os
import sys
import threading
import urllib2
from urlparse import urlparse

# Mostly stolen from the google tools; sets up python search paths so imports from google.* work:
# This won't work on Windows.
DIR_PATH = os.path.abspath(os.path.dirname(os.path.realpath(commands.getoutput("which appcfg.py"))))
SCRIPT_DIR = os.path.join(DIR_PATH, 'google', 'appengine', 'tools')
EXTRA_PATHS = [
  DIR_PATH,
  os.path.join(DIR_PATH, 'lib', 'antlr3'),
  os.path.join(DIR_PATH, 'lib', 'django'),
  os.path.join(DIR_PATH, 'lib', 'webob'),
  os.path.join(DIR_PATH, 'lib', 'yaml', 'lib'),
]
sys.path = EXTRA_PATHS + sys.path

from google.appengine.ext import db
from google.appengine.ext.remote_api import remote_api_stub
from google.appengine.ext.remote_api import throttle as remote_api_throttle
from google.appengine.tools import appengine_rpc

from BTL import BTFailure
from BTL.bencode import bencode, bdecode

from progressbar.progressbar import ProgressBar, Bar, Percentage, ETA, FileTransferSpeed

from create_pieces import *
from model import *

DEFAULT_PIECE_LENGTH = 1<<18  # 256K

def create_torrent_from_dir(path, print_progress):
    if path.endswith(os.sep): path=path[0:-1]
    path_dirs = path.split(os.sep)
    name = path_dirs[-1]
    torrent = { 'info' :
                    { 'name' : name,
                      'piece length' : DEFAULT_PIECE_LENGTH,
                      'pieces' : '',
                      'files' : [ ]
                    }
              }
    file_paths = [ ]
    total_length = 0
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.startswith("."): continue   # Ignore .files
            full_path = os.path.join(root, f)
            file_paths.append(full_path)
            length = os.path.getsize(full_path)
            total_length += length
            full_path_dirs = full_path.split(os.sep)
            torrent['info']['files'].append( { 'length' : length, 'path' : full_path_dirs[len(path_dirs):] } )
        for d in dirs:
            if d.startswith("."): dirs.remove(d)  # Ignore .directories

    # Compute piece hashes:
    widgets = ['Computing piece hashes: ', Bar()]
    if print_progress:
        pbar = ProgressBar(widgets=widgets, maxval=total_length)
        pbar.start()
    nbytes_hashed = 0
    for data in create_pieces_files(torrent['info']['piece length'], '', file_paths):
        torrent['info']['pieces'] += hashlib.sha1(data).digest()
        nbytes_hashed += len(data)
        if print_progress: pbar.update(nbytes_hashed)
    if print_progress: pbar.finish()

    return (torrent, total_length)

def create_torrent_from_file(path, print_progress):
    total_length = os.path.getsize(path)
    torrent = { 'info' :
                    { 'name' : os.path.basename(path),
                      'piece length' : DEFAULT_PIECE_LENGTH,
                      'pieces' : '',
                      'length' : total_length
                    }
              }
    # Compute piece hashes:
    widgets = ['Computing piece hashes: ', Bar()]
    if print_progress:
        pbar = ProgressBar(widgets=widgets, maxval=total_length)
        pbar.start()
    nbytes_hashed = 0
    for data in create_pieces_file(torrent['info']['piece length'], '', path):
        torrent['info']['pieces'] += hashlib.sha1(data).digest()
        nbytes_hashed += len(data)
        if print_progress: pbar.update(nbytes_hashed)
    if print_progress: pbar.finish()

    return (torrent, total_length)

def read_torrent(path):
    fd = open(path, "rb")
    torrent_data = fd.read()
    fd.close()
    torrent = bdecode(torrent_data)

    # TODO: validate?  make sure all files are readable, and all
    # piece hashes compute OK?  And rebuild torrent if piece length is
    # larger than google app engine request size? (10MB)
    # And/or maybe generate a torrent on the
    # fly that omits any missing files.

    if 'files' in torrent['info']:
        return (torrent, sum(f['length'] for f in torrent['info']['files']))
    return (torrent, torrent['info']['length'])

def main():
    parser = optparse.OptionParser(usage="%prog [options] file_or_directory_or_.torrent")
    parser.add_option("--url", dest="url", default="http://localhost:8080/remote_api",
                  help="URL for app engine application (default: %default)")
    parser.add_option("-e", "--email", dest="email", default=None,
                  help="Username to use. Will prompt if omitted (not needed if uploading to localhost)")
    parser.add_option("-p", "--password", dest="password", default=None,
                  help="Password to use.  Will prompt if omitted (not needed if uploading to localhost)")
    parser.add_option("--comment", dest="comment", default=None,
                  help="Add a comment to the .torrent file")
    parser.add_option("--delete", action="store_true", dest="delete", default=False,
                  help="Delete torrent by infohash (usage: --delete HASH)")
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False,
                  help="Print progress bar, infohash and final download URL")
    parser.add_option("--progress", action="store_true", dest="progress", default=False,
                  help="Print progress bar during upload")
    parser.add_option("--info", action="store_true", dest="info", default=False,
                  help="Print info hash (hex-encoded 20-byte SHA1 string) to stdout")
    (options, args) = parser.parse_args()
    print_progress = (options.verbose or options.progress)
    
    if len(args) == 0:
        parser.print_help()
        sys.exit(1)

    url = urlparse(options.url, "http", False)  # Returns 6-tuple: scheme://netloc/path;parameters?query#fragment
    is_dev_server = url[1].startswith('localhost:')

    if is_dev_server:
        authorize = lambda: ('test@example.com', 'password') # Dev server doesn't care what email/password used
    elif options.email is not None and options.password is not None:
        authorize = lambda: (options.email, options.password)
    elif options.email is not None:
        authorize = lambda: (options.email, getpass.getpass('Password: '))
    else:
        authorize = lambda: (raw_input('Google Account Email: '), getpass.getpass('Password: '))

    
    if is_dev_server:
        # Don't need to throttle if talking to dev server
        rpc_server_factory = appengine_rpc.HttpRpcServer
    else:
        # TODO: allow no throttling with command-line option
        full_throttle = remote_api_throttle.DefaultThrottle(multiplier=1.0)
        full_throttle.Register(threading.currentThread())  # only single thread, but have to register anyway
        rpc_server_factory = (
            remote_api_throttle.ThrottledHttpRpcServerFactory(full_throttle))

    remote_api_stub.ConfigureRemoteDatastore(
        None, # Ask server for application id
        url[2],
        authorize,
        servername=url[1],
        rpc_server_factory=rpc_server_factory,
        secure=(url[0] == "https")
        )
    if not is_dev_server:
        remote_api_throttle.ThrottleRemoteDatastore(full_throttle)

    for path in args:
        if path.endswith(os.sep): path=path[0:-1]

        # --delete:
        if options.delete:
            torrent = Torrent.get_by_hash(path[0:40]) # path can be {hash} or {hash.torrent}
            if torrent is None:
                print("ERROR deleting %s-- not found."%path)
                sys.exit(1)
            torrent.delete()
            if options.verbose:
                print("Deleted torrent "+path)
            sys.exit(0)

        # Three upload cases: directory, .torrent, or single-file:
        if os.path.isdir(path):
            (torrent_dict, data_size) = create_torrent_from_dir(path, print_progress)
        else:
            # Right.  If this is a torrent-sized file (small, less than 100,000 bytes), then
            # just try to read it:
            torrent_dict = None
            if os.path.getsize(path) < 100000:
                try:
                    (torrent_dict, data_size) = read_torrent(path)
                except BTFailure:
                    if path.endswith(".torrent"): raise
                    torrent_dict = None    # Fall through to the single-file case
            if torrent_dict is None:
                (torrent_dict, data_size) = create_torrent_from_file(path, print_progress)
    
        if options.comment is not None:
            torrent_dict['comment'] = options.comment
        torrent = Torrent.create_from_dict(torrent_dict)

        parent_dir = os.path.dirname(path)
        piece_length = torrent.info['piece length']
        if 'files' in torrent.info:  # Multi-file torrent
            file_names = [ 
                os.path.join(parent_dir, torrent.info['name'], *f['path']) for f in torrent.info['files'] ]
        else: # Single-file torrent
            file_names = [ os.path.join(parent_dir, torrent.info['name']) ]

        torrent.put()

        widgets = ['Uploading data: ', Bar(), ' ', Percentage(), ' ', ETA()]
        if print_progress:
            pbar = ProgressBar(widgets=widgets, maxval=data_size)
            pbar.start()
        nbytes_transferred = 0
        for piece_number, piece in enumerate(create_pieces_files(piece_length, "", file_names)):
            torrent.put_piece(piece_number, piece)
            nbytes_transferred += len(piece)
            if print_progress: pbar.update(nbytes_transferred)

        torrent.complete = True
        torrent.put()
            
        if print_progress: pbar.finish()
    
        if options.info:
            print(torrent.key().name())
        if options.verbose:
            info_hash = torrent.key().name()
            print("Successfully uploaded %s; info_hash is: %s"%(path, info_hash))
            print("Torrent URL is probably: %s://%s/seed/%s.torrent"%(url[0], url[1], info_hash))

if __name__ == '__main__':
    logging.basicConfig(level=logging.WARNING)
    main()
